import ChatGLM

from langchain_community.document_loaders import PyPDFLoader

from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.embeddings import JinaEmbeddings
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import HuggingFacePipeline
import gradio
from langchain_core.runnables import RunnableParallel, RunnablePassthrough


# llm = ChatGLM.ChatGLM_LLM()

llm = HuggingFacePipeline.from_model_id(
    model_id="THUDM/chatglm3-6b",
      # model_id="google/gemma-7b",
    task="text-generation",
    device=0,
    model_kwargs={"trust_remote_code":True},
    pipeline_kwargs={"max_new_tokens": 5000},
)

loader = PyPDFLoader("西游记.pdf")
documents = loader.load_and_split()

embeddings = HuggingFaceEmbeddings()

# embeddings = JinaEmbeddings(
#     jina_api_key="jina_c5d02a61c97d4d79b88234362726e94aVLMTvF38wvrElYqpGYSxFtC5Ifhj", model_name="jina-embeddings-v2-base-zh"
# )

# 第一次存入本地
# vectorstore = Chroma.from_documents(documents, embeddings,persist_directory="./chroma_db")

# 从本地加载
vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)

retriever = vectorstore.as_retriever()
template = """Answer the question based only on the following context,if can not ，please just say： I do not know，
please think step by step:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
llm = ChatGLM.ChatGLM_LLM()
output_parser = StrOutputParser()
setup_and_retrieval = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
)
chain = setup_and_retrieval | prompt | llm | output_parser
# print(chain.invoke("介绍下红楼梦"))
# print(chain.invoke("第二十二回和第三十回各讲了什么"))
# print(chain.invoke("总结下这本书"))
# print(chain.invoke("孙悟空在这本书里面总共打了多少个妖怪"))

def greet(name):
    response = chain.invoke(name)
    return response
demo = gradio.Interface(fn=greet, inputs="text", outputs="text")
demo.launch() 
